package com.intmall.flink.state;

import com.intmall.flink.operator.ClickSource;
import com.intmall.flink.operator.Event;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.time.Duration;

public class StateTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                        .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                            @Override
                            public long extractTimestamp(Event element, long recordTimestamp) {
                                return element.getTimestamp();
                            }
                        }));

        stream.keyBy(data -> data.getUser())
                .flatMap(new MyFlatMap())
                .print();
        env.execute();
    }

    public static class MyFlatMap extends RichFlatMapFunction<Event, String> {

        private ValueState<Event> myValueState;
        private ListState<Event> myListState;
        private MapState<String, Long> myMapState;
        private ReducingState<Event> myReducingState;
        private AggregatingState<Event, String> myAggregatingState;

        @Override
        public void open(Configuration parameters) throws Exception {
            ValueStateDescriptor<Event> valueStateDescriptor = new ValueStateDescriptor<>("my-state", Types.POJO(Event.class));
            myValueState = getRuntimeContext().getState(valueStateDescriptor);

            // 配置状态的TTL
            StateTtlConfig ttlConfig = StateTtlConfig.newBuilder(Time.seconds(10))
                    .setUpdateType(StateTtlConfig.UpdateType.OnCreateAndWrite)
                    .setStateVisibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
                    .build();
            valueStateDescriptor.enableTimeToLive(ttlConfig);

            myListState = getRuntimeContext().getListState(new ListStateDescriptor<Event>("list-state",
                    Types.POJO(Event.class)));
            myMapState = getRuntimeContext().getMapState(new MapStateDescriptor<String, Long>("map-state",
                    Types.STRING, Types.LONG));
            myReducingState = getRuntimeContext().getReducingState(new ReducingStateDescriptor<Event>("reducing-state",
                    new ReduceFunction<Event>() {
                        @Override
                        public Event reduce(Event value1, Event value2) throws Exception {
                            return new Event(value1.getUser(), value1.getUrl(), value2.getTimestamp());
                        }
                    }, Types.POJO(Event.class)));
            myAggregatingState = getRuntimeContext().getAggregatingState(new AggregatingStateDescriptor<Event, Long, String>("aggregating-state",
                    new AggregateFunction<Event, Long, String>() {
                        @Override
                        public Long createAccumulator() {
                            return 0L;
                        }

                        @Override
                        public Long add(Event value, Long accumulator) {
                            return accumulator + 1;
                        }

                        @Override
                        public String getResult(Long accumulator) {
                            return "count: " + accumulator;
                        }

                        @Override
                        public Long merge(Long a, Long b) {
                            return a + b;
                        }
                    }, Types.LONG));
        }

        @Override
        public void flatMap(Event value, Collector<String> out) throws Exception {
            System.out.println("before: " + myValueState.value());
            myValueState.update(value);
            System.out.println("after: " + myValueState.value());

            myListState.add(value);

            myMapState.put(value.getUser(), myMapState.contains(value.getUser()) ? myMapState.get(value.getUser()) + 1L : 1L);

            myReducingState.add(value);

            myAggregatingState.add(value);
            System.out.println("aggregating: " + myAggregatingState.get());
        }
    }
}
